#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
import can
from typing import Dict
from threading import Lock
import struct
import json
import time
import math
from dummyx_interface.srv import InitUsb2Can
from dummyx_interface.srv import WriteUsb2Can
from dummyx_interface.srv import ReadUsb2Can
from sensor_msgs.msg import JointState

# CAN Command Definitions
CMD_ID_MOTOR_DISABLE = 0x00
CMD_ID_MOTOR_ENABLE = 0x01
CMD_ID_GET_STATUS = 0x0D
CMD_ID_STATUSWORD_REPORT = 0x0E
CMD_ID_GET_VALUE1 = 0x0F
CMD_GET_CONFIG = 0x12
CMD_ID_GET_SAVED_POSITION = 24
CMD_ID_SET_POSITION = 0x04
CMD_ID_SET_VELOCITY = 0x03
CMD_SET_CONFIG = 0x11
CMD_SEND_SYNC = 21
PROFILE_SYNC = 17
CONTROL_MODE_INDEX = 11

class Motor:
    def __init__(self, bus: can.Bus, node_id: int, reduction: float):
        self.bus = bus
        self.node_id = node_id
        self.reduction = reduction
        self.position = 0.0
        self.saved_position = 0.0
        self.enabled = False
        self.lock = Lock()
        self.last_sent_position = None  # Track last sent position
    
    def build_can_id(self, dir_bit: int, cmd_id: int) -> int:
        return (dir_bit << 10) | ((self.node_id & 0x1F) << 5) | (cmd_id & 0x1F)

    def rads_to_rps(self, rad_per_sec):
        return rad_per_sec / (2 * math.pi)

    def send_position(self, position, velocity):
        # Check if position change is significant enough (0.01 degree threshold)
        if (self.last_sent_position is not None and 
            abs(position - self.last_sent_position) < 0.0001):
            return  # Skip sending if change is too small
        print(f"Sent SET_POSITION {self.node_id}: {position:.3f} turns Velocity {velocity:.3f} Velocity {abs(self.rads_to_rps(velocity)):.3f}") 
        # return    

        tx_id = self.build_can_id(dir_bit=0, cmd_id=CMD_ID_SET_VELOCITY)

        data = struct.pack("<f", abs(self.rads_to_rps(velocity)))  # float32类型数据打包
        msg = can.Message(
            arbitration_id=tx_id,
            data=data,
            is_extended_id=False
        )
        
        with self.lock:
            self.bus.send(msg)

        tx_id = self.build_can_id(dir_bit=0, cmd_id=CMD_ID_SET_POSITION)

        data = struct.pack("<f", position)  # float32类型数据打包
        msg = can.Message(
            arbitration_id=tx_id,
            data=data,
            is_extended_id=False
        )
        
        with self.lock:
            self.bus.send(msg)
            self.last_sent_position = position  # Update last sent position
        # print(f"Sent SET_POSITION: {position} turns")

    def enable(self):

        tx_id = self.build_can_id(dir_bit=0, cmd_id=CMD_SET_CONFIG)
        # change to mode 2
        data = struct.pack("<II", CONTROL_MODE_INDEX, 2)
        msg = can.Message(
            arbitration_id=tx_id,
            data=data,
            is_extended_id=False
        )
        with self.lock:
            self.bus.send(msg)
        
        time.sleep(0.1)

        tx_id = self.build_can_id(dir_bit=0, cmd_id=CMD_SET_CONFIG)
        data = struct.pack("<II", PROFILE_SYNC, 1)
        msg = can.Message(
            arbitration_id=tx_id,
            data=data,
            is_extended_id=False
        )
        with self.lock:
            self.bus.send(msg)
        
        time.sleep(0.1)

        tx_id = self.build_can_id(dir_bit=0, cmd_id=CMD_ID_MOTOR_ENABLE)
        msg = can.Message(
            arbitration_id=tx_id,
            data=[],
            is_extended_id=False
        )
        with self.lock:
            self.bus.send(msg)
            self.enabled = True
    
    def disable(self):
        tx_id = self.build_can_id(dir_bit=0, cmd_id=CMD_ID_MOTOR_DISABLE)
        msg = can.Message(
            arbitration_id=tx_id,
            data=[],
            is_extended_id=False
        )
        with self.lock:
            self.bus.send(msg)
            self.enabled = False
    
    def update_status(self, msg: can.Message):
        if len(msg.data) >= 8:
            with self.lock:
                self.position = struct.unpack("<f", msg.data[4:8])[0]
    
    def update_saved_position(self, msg: can.Message):
        if len(msg.data) >= 4:
            with self.lock:
                self.saved_position = struct.unpack("<f", msg.data[0:4])[0]

class MotorController:
    def __init__(self, interface: str, channel: str, motor_config: Dict[int, float]):
        self.bus = can.Bus(interface=interface, channel=channel, receive_own_messages=True)
        self.motors: Dict[int, Motor] = {}
        self.message_lock = Lock()
        
        # Initialize all motors
        for node_id, reduction in motor_config.items():
            self.motors[node_id] = Motor(self.bus, node_id, reduction)
        
        # Start CAN message listener
        self.notifier = can.Notifier(self.bus, [self.on_message_received])
    
    def build_can_bc(self, dir_bit: int, cmd_id: int) -> int:
        return (dir_bit << 10) | ((0x1f & 0x1F) << 5) | (cmd_id & 0x1F)

    def send_sync(self):
        tx_id = self.build_can_bc(dir_bit=0, cmd_id=CMD_SEND_SYNC)
        msg = can.Message(
            arbitration_id=tx_id,
            data=[],
            is_extended_id=False
        )
        with self.message_lock:
            self.bus.send(msg)

    def enable_all_motors(self):
        """Enable all initialized motors"""
        for motor in self.motors.values():
            motor.enable()
            time.sleep(0.1)  # Small delay between commands
    
    def disable_all_motors(self):
        """Disable all motors"""
        for motor in self.motors.values():
            motor.disable()
            time.sleep(0.1)
    
    def on_message_received(self, msg):
        dir_bit = (msg.arbitration_id >> 10) & 0x1
        node_id = (msg.arbitration_id >> 5) & 0x1F
        cmd_id = msg.arbitration_id & 0x1F
        
        if dir_bit == 1 and node_id in self.motors:
            if cmd_id == CMD_ID_STATUSWORD_REPORT:
                self.motors[node_id].update_status(msg)
            elif cmd_id == CMD_ID_GET_STATUS:
                self.motors[node_id].update_status(msg)
            elif cmd_id == CMD_ID_GET_VALUE1:
                pass  # Simplified - no status update
            elif cmd_id == CMD_GET_CONFIG:
                pass  # Simplified - no config update
            elif cmd_id == CMD_ID_GET_SAVED_POSITION:
                self.motors[node_id].update_saved_position(msg)
    
    def shutdown(self):
        self.notifier.stop()
        self.bus.shutdown()

class USB2CANNode(Node):
    def __init__(self):
        super().__init__('usb2can_node')
        
        # Declare parameters
        self.declare_parameter('interface', 'socketcan')
        self.declare_parameter('channel', 'can0')
        self.declare_parameter('motor_config_json', '{"1":50.0,"2":50.0,"3":50.0,"4":50.0,"5":50.0,"6":50.0}')
        self.declare_parameter('joint_names', ['Joint1', 'Joint2', 'Joint3', 'Joint4', 'Joint5', 'Joint6'])

        self.srv = self.create_service(InitUsb2Can, 'init_usb2can', self.init_usb2can_callback)
        self.srv = self.create_service(WriteUsb2Can, 'write_usb2can', self.write_usb2can_callback)
        self.srv = self.create_service(ReadUsb2Can, 'read_usb2can', self.read_usb2can_callback)
        
        # Get parameter values
        interface = self.get_parameter('interface').value
        channel = self.get_parameter('channel').value
        motor_config_json = self.get_parameter('motor_config_json').value
        self.joint_names = self.get_parameter('joint_names').value
        
        try:
            motor_config = {int(k): float(v) for k, v in json.loads(motor_config_json).items()}
        except (json.JSONDecodeError, ValueError) as e:
            self.get_logger().error(f"Invalid motor_config_json format: {e}")
            motor_config = {1: 50.0, 2: 50.0, 3: 50.0, 4: 50.0, 5: 50.0, 6: 50.0}
        
        self.get_logger().info(f"Initializing USB2CAN with interface: {interface}, channel: {channel}")
        self.get_logger().info(f"Motor configuration: {motor_config}")
        
        try:
            # Initialize motor controller with all motors
            self.controller = MotorController(
                interface=interface,
                channel=channel,
                motor_config=motor_config
            )
            self.get_logger().info("Successfully initialized all motors")
            
            # Enable all motors
            self.controller.enable_all_motors()
            self.get_logger().info("Enabled all motors")
            
            # Create subscriber for joint states
            self.joint_state_sub = self.create_subscription(
                JointState,
                'joint_states',
                self.joint_state_callback,
                10
            )
            self.get_logger().info("Subscribed to joint_states topic")
            
        except Exception as e:
            self.get_logger().error(f"Failed to initialize USB2CAN: {str(e)}")
            raise
    
    def joint_state_callback(self, msg: JointState):
        # Check if we have all expected joints
        if not all(name in msg.name for name in self.joint_names):
            self.get_logger().warn("Received joint states don't contain all expected joints")
            return
        
        # Convert joint positions from radians to degrees and send to motors
        for i, joint_name in enumerate(self.joint_names):
            try:
                joint_index = msg.name.index(joint_name)
                joint_position_rad = msg.position[joint_index]
                joint_velocity_rad = msg.velocity[joint_index]
                joint_position_deg = math.degrees(joint_position_rad)
                
                # Motor IDs are 1-6 corresponding to joints 1-6
                motor_id = i + 1
                if motor_id in self.controller.motors:
                    # Calculate motor turns: degrees / 360 * reduction
                    motor_turns = joint_position_deg / 360 * self.controller.motors[motor_id].reduction
                    if motor_id == 2:
                        self.controller.motors[motor_id].send_position(motor_turns * -1, joint_velocity_rad * self.controller.motors[motor_id].reduction)
                    else:
                        self.controller.motors[motor_id].send_position(motor_turns, joint_velocity_rad * self.controller.motors[motor_id].reduction)
                #everything good send out sync command
                self.controller.send_sync()
            except (ValueError, IndexError) as e:
                self.get_logger().error(f"Error processing joint state for {joint_name}: {str(e)}")

    def init_usb2can_callback(self, request, response):
        self.get_logger().info("Initializing USB2CAN callback")
        if request.action == 'stop':
            self.controller.disable_all_motors()
        elif request.action == 'start':
            self.controller.enable_all_motors()
        response.success = True
        return response

    def write_usb2can_callback(self, request, response):
        if self.controller is None:
            response.success = False
            return response
        response.success = True  
        return response

    def read_usb2can_callback(self, request, response):
        response.pos_commands = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        response.vel_commands = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        response.success = True  
        return response

    def on_shutdown(self):
        self.get_logger().info("Shutting down USB2CAN controller")
        if hasattr(self, 'controller'):
            self.controller.disable_all_motors()
            self.controller.shutdown()

def main(args=None):
    rclpy.init(args=args)
    node = None
    
    try:
        node = USB2CANNode()
        rclpy.spin(node)
    except Exception as e:
        if node is not None:
            node.get_logger().error(f"Error in USB2CAN node: {str(e)}")
        else:
            print(f"Error creating USB2CAN node: {str(e)}")
    finally:
        if node is not None:
            node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()